from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langgraph.prebuilt import create_react_agent

from langchain_openai import ChatOpenAI

from db_utils import company_query

from sshcode.db_agent.model_utils import getLLM

_table_system_template = """
您是一个注重SQL专家
给定一个查询目标，生成一个语法正确的SQLite查询语句，然后调用查询工具将查询的结果返回

工具:
{tools_name}

"""

_table_name_template = """
查询目标: 获得表{table_name}表的所有信息,包括表名和字段名。

"""

class Info:
    def __init__(self,llm):
        _tools = [company_query]
        _prompt = ChatPromptTemplate.from_messages([
            ("system",_table_system_template),
            ("human",_table_name_template)
        ])
        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools ]))

        # 工具代理
        _llm_with_tool_agent = create_react_agent(llm,tools=_tools)

        self._parser = StrOutputParser()
        self._chain = _prompt|_llm_with_tool_agent
    # # TypeError: 'Info' object is not callable
    # def _call__(self,state):
    def __call__(self,state):
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]
        return self._parser.invoke(_messages[-1])


if __name__ == '__main__':
    _llm = getLLM()
    # TypeError: 'Info' object is not callable
    _info = Info(_llm)
    _rt = _info({"table_name": "employees"})
    print(_rt)